package bosc.cn.tunnel;

import bosc.cn.RustSignReqDTO;
import bosc.cn.RustSignResDTO;
import bosc.cn.RustVerifyReqDTO;
import bosc.cn.RustVerifyResDTO;
import bosc.cn.util.JsonUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.client.WebSocketClient;

import javax.annotation.PostConstruct;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @author tianshq
 * @date 2025/9/8 23:22
 * @description SignServiceTunnelClient
 * sign-service（外网服务器）
 * 通过隧道调用 rust-service：
 */
@Slf4j
@Component
public class SignServiceTunnelClient {

    public static final String SIGN = "sign";
    public static final String VERIFY = "verify";
    public static final String REGEX = "@!@";
    public static final String RUST_SERVICE = "rust-service";
    public static final String WEBSOCKET_REGEX = "|";

    @Value("${tunnel.server.url}")
    private String tunnelServerUrl;

    @Autowired
    private WebSocketClient webSocketClient;

    // 使用连接池管理多个 WebSocketSession
    private final Map<Integer, WebSocketSession> sessionPool = new ConcurrentHashMap<>();
    // 生成请求ID的计数器
    private final AtomicLong requestIdCounter = new AtomicLong(0);
    // 使用 Map 管理每个请求的 responseFuture，通过请求ID区分
    private final Map<String, CompletableFuture<String>> responseFutures = new ConcurrentHashMap<>();


    @PostConstruct
    public void connect() throws Exception {
        // 初始化多个连接，例如初始化3个连接
        for (int i = 0; i < 3; i++) {
            WebSocketSession session = webSocketClient.doHandshake(
                    new TunnelWebSocketHandler(this),
                    tunnelServerUrl
            ).get();
            sessionPool.put(i, session);
        }
    }

    // 根据索引选择连接
    private WebSocketSession getSession(int index) {
        return sessionPool.get(index % sessionPool.size());
    }

    // 生成唯一请求ID
    private String generateRequestId() {
        return "req_" + requestIdCounter.incrementAndGet();
    }

    public CompletableFuture<RustSignResDTO> callRustSign(RustSignReqDTO request, int realIndex) throws Exception {
        // 生成唯一请求ID
        String requestId = generateRequestId();
        // 设置局部的响应等待，不能共用，不然容易响应混乱
        CompletableFuture<String> responseFuture = new CompletableFuture<>();
        responseFutures.put(requestId, responseFuture);

        // 发送请求到隧道服务器
        // 格式: clientId前缀|methodName@!@oriRequestId@!@JSON
        StringBuilder tunnelMessage = new StringBuilder();
        tunnelMessage.append(RUST_SERVICE).append("-").append(realIndex)
                .append(WEBSOCKET_REGEX).append(SIGN).append(REGEX)
                .append(requestId).append(REGEX)
                .append(JsonUtils.toJson(request));
        // 选择连接
        WebSocketSession tunnelSession = getSession(realIndex);
        try {

            tunnelSession.sendMessage(new TextMessage(tunnelMessage.toString()));

            // 返回 CompletableFuture，不阻塞等待
            return responseFuture.thenApply(response -> {
                log.info("Rust-service-{}，返回数据：{}", realIndex, response);
                String[] req_json = response.split(REGEX);
                return JsonUtils.fromJson(req_json[1], RustSignResDTO.class);
            }).whenComplete((result, throwable) -> {
                // 清理资源
                responseFutures.remove(requestId);
                if (throwable != null) {
                    log.error("处理响应时发生错误", throwable);
                }
            });
        } catch (Exception e) {
            responseFutures.remove(requestId);
            CompletableFuture<RustSignResDTO> failedFuture = new CompletableFuture<>();
            failedFuture.completeExceptionally(e);
            return failedFuture;
        }
    }

    // 请求rust完成验证
    public RustVerifyResDTO callRustVerify(RustVerifyReqDTO request, int realIndex) throws Exception {
        // 生成唯一请求ID
        String requestId = generateRequestId();
        // 设置响应等待
        CompletableFuture<String> responseFuture = new CompletableFuture<>();
        responseFutures.put(requestId, responseFuture);


        // 发送请求到隧道服务器
        // 格式: clientId前缀|methodName@!@oriRequestId@!@JSON
        StringBuilder tunnelMessage = new StringBuilder();
        tunnelMessage.append(RUST_SERVICE).append("-").append(realIndex)
                .append(WEBSOCKET_REGEX).append(VERIFY).append(REGEX)
                .append(requestId).append(REGEX)
                .append(JsonUtils.toJson(request));
        // 使用第一个连接（或根据需要选择）
        WebSocketSession tunnelSession = getSession(realIndex);
        try {
            // 使用同步块保护 WebSocketSession 的访问
            //synchronized (lock) {
            tunnelSession.sendMessage(new TextMessage(tunnelMessage.toString()));
            //}


            // 等待响应
            String response = responseFuture.get(60, TimeUnit.SECONDS);

            // 与RUST约定一致格式
            // methodName@!@JSON
            log.info("Rust-service返回数据：{}", response);
            String[] req_json = response.split(REGEX);


            // 解析响应
            return JsonUtils.fromJson(req_json[1], RustVerifyResDTO.class);
        } finally {
            // 清理资源
            responseFutures.remove(requestId);
        }
    }

    // 供 WebSocketHandler 调用设置响应
    public void setResponse(String response, String requestId) {
        CompletableFuture<String> responseFuture = responseFutures.get(requestId);
        if (responseFuture != null && !responseFuture.isDone()) {
            responseFuture.complete(response);
        }
    }
}